perf: optimize date_trunc (>2x faster)#4915
Conversation
date_trunc (>2x faster)
mbutrovich
left a comment
There was a problem hiding this comment.
First pass, thanks @andygrove!
| fn date_trunc_fn_for_format(format: &str) -> Result<DateTruncFn, SparkError> { | ||
| let key: Cow<str> = if format.is_ascii() { | ||
| Cow::Borrowed(format) | ||
| } else { | ||
| Cow::Owned(format.to_uppercase()) | ||
| }; | ||
| DATE_TRUNC_FORMATS | ||
| .iter() | ||
| .find(|(name, _)| name.eq_ignore_ascii_case(&key)) | ||
| .map(|(_, trunc_fn)| *trunc_fn) | ||
| .ok_or_else(|| { | ||
| SparkError::Internal(format!( | ||
| "Unsupported format: {format:?} for function 'date_trunc'" | ||
| )) | ||
| }) | ||
| } |
There was a problem hiding this comment.
let key: Cow<str> = if format.is_ascii() {
Cow::Borrowed(format)
} else {
Cow::Owned(format.to_uppercase())
};
DATE_TRUNC_FORMATS
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case(&key))Every entry in DATE_TRUNC_FORMATS is ASCII. eq_ignore_ascii_case only folds ASCII bytes, so a non-ASCII key can never equal an ASCII key regardless of whether it was to_uppercased first. The else branch allocates a String and then cannot change the outcome: a non-ASCII format always falls through to the error. The comment claiming non-ASCII "still needs full Unicode case folding to fold the same way Spark does" is incorrect for this table, because Spark also only matches ASCII literals after folding, so a non-ASCII input is invalid in Spark too.
Suggested change: drop the Cow and the is_ascii() split entirely and match on the borrowed &str:
fn date_trunc_fn_for_format(format: &str) -> Result<DateTruncFn, SparkError> {
DATE_TRUNC_FORMATS
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case(format))
.map(|(_, trunc_fn)| *trunc_fn)
.ok_or_else(|| {
SparkError::Internal(format!(
"Unsupported format: {format:?} for function 'date_trunc'"
))
})
}This removes the use std::borrow::Cow; import, removes an allocation on the non-ASCII path, and is exactly as Spark-compatible.
There was a problem hiding this comment.
Dropped the Cow / is_ascii() split. Since every entry in DATE_TRUNC_FORMATS is ASCII and eq_ignore_ascii_case only folds ASCII, a non-ASCII input can never match either way — the fallback allocation was pure cost.
There was a problem hiding this comment.
native/spark-expr/src/kernels/temporal.rs:591 (ntz_trunc_fn_for_format uses format.to_uppercase().as_str()) and 832 (the tz array-format helper still calls $formats.value(index).to_uppercase().as_str() per row). These are the timestamp_trunc siblings of the code this PR just optimized, and they carry the identical per-row String allocation the PR is removing. Leaving them behind means the "resolve once into a table plus memo" pattern this PR introduces is applied to one of three structurally identical loops, so the improvement is narrower than the general problem.
Suggested change: either extend this PR to build a TIMESTAMP_TRUNC_FORMATS table plus the same last-format memo for the NTZ and tz array paths, or file a follow-up issue and reference it in this PR so the remaining allocations are tracked. I would fold at least the NTZ path in here since it is a direct analog of DATE_TRUNC_FORMATS.
There was a problem hiding this comment.
Extended the table pattern to ntz_trunc_fn_for_format and added a monomorphized-for-DateTime<Tz> sibling tz_trunc_fn_for_format. Both the tz array-format helper and the scalar-format tz path now resolve via eq_ignore_ascii_case against a static table, so the per-row and per-invocation to_uppercase() allocations are gone.
Applies the DATE_TRUNC_FORMATS table pattern to both timestamp_trunc paths. ntz_trunc_fn_for_format and tz_trunc_fn_for_format resolve the format via eq_ignore_ascii_case against ASCII tables, eliminating the per-invocation and per-row to_uppercase() String allocations from the scalar-format tz path and the array-format helper. Also drops the unused Cow / is_ascii split from date_trunc_fn_for_format since eq_ignore_ascii_case only folds ASCII and every table entry is ASCII, so a non-ASCII input cannot match either way.
Which issue does this PR close?
N/A
Rationale for this change
Optimize existing expression.
What changes are included in this PR?
Resolve the date_trunc format keyword without allocating an uppercased String per row (ASCII case-insensitive match plus a last-format memo), removing a heap allocation from every row of the array-format path.
How are these changes tested?
Existing tests.
Benchmark (criterion):
Full criterion output: